# 作者: 肖老师
# 2024年12月19日04时20分16秒
# xxxrlmk@163.com
RED = 0  # 红色
BLACK = 1  # 黑色


class RBNode:
    def __init__(self, value, color=RED):
        self.value = value
        self.color = color
        self.p = None
        self.left = None
        self.right = None


class RedBlackTree:
    def __init__(self):
        self.root: RBNode = None

    def right_rotate(self, y: RBNode):
        x: RBNode = y.left
        # 父亲的处理
        x.p = y.p
        if y.p is None:
            self.root = x  # x作为新的根
        else:
            if y.p.left == y:  # 说明y在父亲左边
                y.p.left = x
            else:
                y.p.right = x
        # x右孩子处理
        y.left = x.right
        if x.right is not None:
            x.right.p = y
        # x与y之间的
        x.right = y
        y.p = x

    def left_rotate(self, x: RBNode):
        if x.right is None:
            return
        y: RBNode = x.right
        if y.left:  # 如果y的左子树不为空,则将y的左子树变为x的右子树
            y.left.p = x  # 将y的左子树的父节点设置为x
            x.right = y.left  # 将y的左子树变为x的右子树
        else:
            x.right = None
        y.p = x.p  # 将y的父节点设置为x的父节点
        if x.p is None:  # 如果x的父节点为空,则将y设置为根节点
            self.root = y
        else:
            if x == x.p.left:  # 如果x是x的父节点的左子树,则将y设置为x的父节点的左子树
                x.p.left = y  # 将y设置为x的父节点的左子树
            else:
                x.p.right = y  # 将y设置为x的父节点的右子树
        y.left = x  # 将x设置为y的左子树
        x.p = y  # 将x的父节点设置为y

    def insert(self, value):
        new_node = RBNode(value)
        # 二叉查找树的插入
        if self.root is None:  # 树根为空
            self.root = new_node
        else:
            x: RBNode = self.root
            # y不可能为空
            while x:
                y: RBNode = x
                if x.value > new_node.value:
                    x = x.left
                else:
                    x = x.right
            # 判断new_node是放在y的左边还是右边
            new_node.p = y
            if y.value < new_node.value:
                y.right = new_node
            else:
                y.left = new_node
        self.insert_fixup(new_node)

    def insert_fixup(self, node):
        parent: RBNode = node.p
        # 父亲存在而且是红色
        while parent and parent.color == RED:
            # 能进到此处，说明爷爷一定存在，不为空，因为原本的树就是一颗红黑树
            gparent: RBNode = parent.p
            if parent == gparent.left:
                # 情形三
                uncle: RBNode = gparent.right
                if uncle and uncle.color == RED:
                    parent.color = BLACK
                    uncle.color = BLACK
                    gparent.color = RED
                    node = gparent
                    parent = node.p
                    continue
                # 情形四
                if parent.right == node:
                    self.left_rotate(parent)
                    node, parent = parent, node
                # 情形五
                self.right_rotate(gparent)
                parent.color = BLACK
                gparent.color = RED
            else:
                # 情形三
                uncle: RBNode = gparent.left
                if uncle and uncle.color == RED:
                    parent.color = BLACK
                    uncle.color = BLACK
                    gparent.color = RED
                    node = gparent
                    parent = node.p
                    continue
                # 情形四
                if parent.left == node:
                    self.right_rotate(parent)
                    node, parent = parent, node
                # 情形五
                self.left_rotate(gparent)
                parent.color = BLACK
                gparent.color = RED

        self.root.color = BLACK  # 情形一


def rbtree_print(node, key, direction):
    if node:
        if direction == 0:  # tree是根节点
            print("%2d(B) is root" % node.value)
        else:  # tree是分支节点
            print("%2d(%s) is %2d's %6s child" % (
                node.value, ("B" if node.color == 1 else "R"), key, ("right" if direction == 1 else "left")))

        rbtree_print(node.left, node.value, -1)
        rbtree_print(node.right, node.value, 1)


if __name__ == '__main__':
    number_list = (7, 4, 1, 8, 5, 2, 9, 6, 3)
    tree = RedBlackTree()
    for i in number_list:
        tree.insert(i)
    rbtree_print(tree.root, tree.root.value, 0)
